import mne
import numpy as np
import matplotlib.pyplot as plt

"""
This script is to preprocess EEG data of patients with depression.
created by JINJing
"""

# 加载脑电数据
# 假设您的数据存储在名为'eeg_data.fif'的文件中
raw = mne.io.read_raw_fif('eeg_data.fif')

# 查看数据的基本信息
print(raw.info)

# 滤波处理
# 例如，使用带通滤波器过滤掉5Hz以下的低频噪声和50Hz的电源干扰
raw.filter(l_freq=5, h_freq=None, notchfreqs=[50])

# 去除伪迹
# 这可能包括眼动、肌电等噪声的去除
# 这里我们使用ICA（独立成分分析）作为示例
ica = mne.preprocessing.ICA(n_components=0.95, random_state=0)
ica.fit(raw)

# 绘制ICA成分的拓扑图，以便选择需要去除的成分
ica.plot_components()

# 假设我们选择去除前n个成分
n_remove = 3  # 根据实际情况调整
ica.exclude_components(range(n_remove))

# 应用ICA并重建数据
raw_clean = ica.apply(raw)

# 分段处理
# 将连续的数据分割成较短的时间段，以便进行后续分析
event_id = {'Depression': 1}  # 根据您的数据设置事件ID
events = mne.find_events(raw_clean, stim_channel='STI 014')  # 根据您的数据设置刺激通道
epochs = mne.Epochs(raw_clean, events, event_id, tmin=-0.5, tmax=0.5, proj=True, preload=True)

# 特征提取
# 这可能包括提取频带功率、时域特征等
# 这里我们简单计算每个epoch的平均功率作为示例
freqs = np.arange(5, 30, 1)  # 5-30Hz范围内的频率
power = mne.time_frequency.psd_multitaper(epochs, fmin=freqs[0], fmax=freqs[-1], faverage=True, tmin=None, tmax=None,
                                          return_onesided=True, n_jobs=1)

# 绘制功率谱图
plt.plot(freqs, power.mean(axis=0).T)
plt.xlabel('Frequency (Hz)')
plt.ylabel('Power')
plt.show()